package test.instrumentation.classfiletest;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;
import javassist.NotFoundException;

public class InstrumentationClass implements ClassFileTransformer {

	private static final String CLASS_TO_INSTRUMENT = "target.ClassToBeInstrumented";

	@Override
	public byte[] transform(final ClassLoader loader, final String className,
			final Class<?> classBeingRedefined,
			final ProtectionDomain protectionDomain,
			final byte[] classfileBuffer) throws IllegalClassFormatException {

		if (!CLASS_TO_INSTRUMENT.equals(classBeingRedefined.getName())) {
			return classfileBuffer;
		}

		try {
			return performClassTransform(loader);
		} catch (final Exception e) {
			throw new RuntimeException(e);
		}
	}

	private byte[] performClassTransform(final ClassLoader loader)
			throws NotFoundException, CannotCompileException, IOException {
		final CtClass ctClass = getCtClass(loader);
		final CtMethod methodToTransform = getMethodToInstrument(ctClass);
		methodToTransform.setBody("return 1;");
		return ctClass.toBytecode();
	}

	private CtClass getCtClass(final ClassLoader loader)
			throws NotFoundException {
		final ClassPool pool = ClassPool.getDefault();
		final LoaderClassPath lcp = new LoaderClassPath(loader);
		pool.appendClassPath(lcp);

		final CtClass ctClass = pool.get(CLASS_TO_INSTRUMENT);
		return ctClass;
	}

	private CtMethod getMethodToInstrument(final CtClass ctClass)
			throws NotFoundException {
		final String methodToBeInstrumented = "shouldReturnOne";
		// see
		// http://java.sun.com/docs/books/jvms/second_edition/html/ClassFile.doc.html
		// for details on descriptors
		final String methodDescriptor = "()I";
		return ctClass.getMethod(methodToBeInstrumented, methodDescriptor);
	}

}
